#include "common.cuh"

__global__ void SGEMM2(_float* A, _float* B, _float* C, int heightA, int widthA, int widthB)
{
	const int HC = heightA;
	const int HB = widthA;
	const int WC = widthB;
	const int tx = threadIdx.x;
	const int ty = threadIdx.y;
	//const int t = threadIdx.y * G2_BLOCK_X + threadIdx.x;
	const int t = ty * G2_BLOCK_X + tx;
	const int ttx= t % G2_C_STRIDE;
	const int tty= t / G2_C_STRIDE;
    const int bbx = blockIdx.x * G2_SMEM_AX;
	const int bby = blockIdx.y * G2_SMEM_AX;
	
	__shared__ _float As[G2_SMEM_AX][G2_SMEM_AY + 1];
    __shared__ _float Bs[G2_SMEM_AY][G2_SMEM_AX + 1];

	// output results of C for this thread block (row major)
	_float a[6];
	_float b[6];
	_float c[36] = {0,0,0,0,0,0,  0,0,0,0,0,0,  0,0,0,0,0,0, 0,0,0,0,0,0, 0,0,0,0,0,0, 0,0,0,0,0,0};

	// initialize pointers to A and B and C for this thread
	const int offsetA = bbx + heightA * tty + ttx;
	const int offsetB = HB * (bby + tty * G2_OUTPUT_X) + ttx;
	//A += bbx + heightA * tty + ttx;
	A += offsetA;
	//B += HB * (bby + tty * G2_OUTPUT_X) + ttx;
	B += offsetB;
	//C += HC * (bby + ttx)  + bbx + tty;	
	C += heightA * (bby + tty)  + bbx + ttx;
	
	const _float* BEnd = B + widthA;
	// iterate over sub-matrices of As and Bs
	while ( B < BEnd) {
//__syncthreads();
		// load shared As and Bs sub matrices from global A and B
		#pragma unroll
		for(int i=0; i < G2_OUTPUT_X; i++) {
			As[i*G2_C_STRIDE + ttx][tty] = A[ i*G2_C_STRIDE];
		}
		__syncthreads();
		// load B after loading A to increase cache hit
		#pragma unroll
		for(int i=0; i < G2_OUTPUT_X; i++) {			
			Bs[ttx][tty*G2_OUTPUT_X + i] = B[i* HB];
		}
		__syncthreads();

		// for each column of As and row of Bs
		#pragma unroll
		for( int k=0; k < 16; ++k) {
			// load the elements this thread is responsible for of As and Bs into registers a[], b[]
			#pragma unroll
			for(int i=0; i < G2_OUTPUT_X; ++i) {
				a[i] = As[i*G2_C_STRIDE + ttx][k];
				b[i] = Bs[k][i*G2_C_STRIDE + tty];
			}
			#pragma unroll
			for(int i=0; i < G2_OUTPUT_X; ++i) {
				#pragma unroll 
				for (int j=0; j < G2_OUTPUT_X; ++j) {
					c[i*G2_OUTPUT_X + j] += a[i] * b[j];
				}
			}
			//__syncthreads();
		}
		// update pointers
		A += heightA * G2_SMEM_AY;
		B += G2_SMEM_AY;

		__syncthreads();
	}

//	__syncthreads();

	// store results in sub-matrix of C for this block
	#pragma unroll
    for( int i=0; i < G2_OUTPUT_X; i++){
 		#pragma unroll
		for(int j=0; j < G2_OUTPUT_X; j++)
			C[j*G2_C_STRIDE] = c[j*G2_OUTPUT_X + i];	
		   
		C += heightA * G2_C_STRIDE;
	}
}